{-# LANGUAGE RecordWildCards #-} 

module Tree where


import Prelude hiding (sequence,  mapM)

import Data.Tree hiding (Tree)
import Data.Maybe
import Data.Traversable (sequence, mapM)

import Control.Monad.State.Lazy hiding (sequence, mapM)

import Graphics.Vty
import Graphics.Vty.Widgets.Base


-- current node is always visible


-- TODO: размеры


-- | Дерево хранится в виде Forest (NodeData a).  
-- Поскольку используется готовое дерево, приходится вручную везде проверять инвариант
-- "если дочерних узлов нет, то узел считается развёрнутым". 
data Tree a
  = Tree
    { selectedAttrFn :: Attr -> Attr
    , nodeAttrFn :: a -> Attr
    , nodeLabelFn :: a -> String
    , selectedIndex :: Int
    , windowStart :: Int
    , treeWidth  :: Maybe Int
    , treeHeight :: Maybe Int
    , nodes :: Forest (NodeData a)
    }

data NodeData a = NodeData {collapsed :: Bool, value :: a} deriving Show


defaultTree tree
  = Tree
    { selectedAttrFn = (`with_back_color` blue)
    , nodeAttrFn = const def_attr
    , nodeLabelFn = id
    , selectedIndex = 0
    , windowStart = 0
    , treeWidth = Nothing
    , treeHeight = Nothing
    , nodes = map (fmap $ NodeData False) tree
    }

  
current t = value $ (concatMap flatten $ map filterVisible $ nodes t) !! selectedIndex t


moveUp t@(Tree {windowStart = ws, selectedIndex = sIx}) 
  | sIx == ws && ws == 0 = t
  | sIx == ws            = t { windowStart = ws-1, selectedIndex = sIx-1 }
  | otherwise            = t { selectedIndex = sIx-1 }


moveDn t@(Tree {windowStart = ws, selectedIndex = sIx}) 
  | sIx == treeSize - 1 = t
  | otherwise           = t { selectedIndex = sIx+1 }
  where
    treeSize = length $ concatMap (flatten . filterVisible) $ nodes t

{- TODO: сделать monad transformer
appendIndex nodes = evalState (mapM (mapNode index) nodes) 0
  where
    index = getIndex >>= \ix -> put (Index $ ix+1) >> return ix
    noIndex = return $ Index (-1)

    mapNode indexProvider (Node n ns) = do
      ix <- indexProvider
      ns' <- mapM (mapNode $ if collapsed n then noIndex else index) ns
      return $ Node (ix,n) ns'
-}

setCollapsed value nodeIx nodes = evalState (mapM setCollapsed' nodes) 0
  where
    setCollapsed' (Node n ns) = do
      ix <- get
      put $ ix + 1
      if ix == nodeIx
          -- проверяем null ns, чтобы не было крестиков около листьев
          then return $ Node (n {collapsed = value && (not.null) ns}) ns
        else if collapsed n || ix > nodeIx
          then return $ Node n ns
        else Node n `fmap` mapM setCollapsed' ns

 
expand t@(Tree {..}) = t {nodes = setCollapsed False selectedIndex nodes}


-- FIXME: работает совсем неправильно
collapse t@(Tree {..}) = t
  { selectedIndex = parentIx
  , nodes = setCollapsed True parentIx nodes
  }
  where
    -- ищем узел, который нужно свернуть. Это может быть:
    --   - выбранный узел, если у него есть вложенные узлы и он сам развёрнут
    --   - иначе, сворачиваем родителя выбранного узла
    parentIx = fst $ foldl getParent1 (0, 0) nodes

    -- для самого верхнего уровня логика немного отличается:
    -- если узел свёрнут или он лист, то ничего не происходит (потому как родительского узла нет)
    getParent1 (p,ix) n = if ix <= selectedIndex then getParent' (ix,ix) n else (p,ix)
    -- у вложенных узлов всегда есть родитель которого можно свернуть
    getParent2 (p,ix) n = if ix <= selectedIndex then getParent' (p,ix) n else (p,ix)

    getParent' (p,ix) (Node n ns)
      -- если выбран лист дерева или свёрнутый узел, нужно сворачивать ветку на которой он висит
      -- если выбран развёрнутый узел дерева, значит его самого и сворачиваем
      | ix == selectedIndex = if null ns || collapsed n then (p,ix+1) else (ix,ix+1)
      -- до выбранного узла ещё не добрались, встретили развёрнутый,
      -- запоминаем его как родительский и заходим внутрь
      | (not.collapsed) n && (not.null) ns = foldl getParent2 (ix,ix+1) ns
      -- встретили лист или свёрнутый узел, проходим мимо
      | otherwise = (p,ix+1)


resize _ h t@(Tree {windowStart = ws, selectedIndex = sIx})
  -- current node is out of view ==> scroll up until it is visible
  | sIx > ws + h - 1 = t { windowStart = sIx - h + 1 }
  -- we have free space after last node ==> scroll down to fill it
  | wl < h           = t { windowStart = max 0 $ ws - (h - wl) }
  | otherwise        = t
  where
    wl = length $ window ws h t


instance Widget (Tree a) where
  -- growHorizontal :: w -> Bool
  growHorizontal _ = True

  -- growVertical :: w -> Bool
  growVertical _ = True

  -- primaryAttribute :: w -> Attr
  primaryAttribute _ = def_attr

  -- withAttribute :: w -> Attr -> w
  withAttribute w _ = w

  -- render :: DisplayRegion -> w -> Image
  render rgn t@(Tree {..})
    = vert_cat
      $  map (renderNode width nodeAttrFn nodeLabelFn) a
      ++ [renderNode width (selectedAttrFn . nodeAttrFn) nodeLabelFn n]
      ++ map (renderNode width nodeAttrFn nodeLabelFn) b
    where
      width  = maybe (fromIntegral $ region_width rgn)  id treeWidth
      height = maybe (fromIntegral $ region_height rgn) id treeHeight
      Tree {windowStart = start} = resize width height t
      (a, n:b) = splitAt (selectedIndex - start) $ window start height t


renderNode width attrFn labelFn (depth, n)
  = string (attrFn $ value n)
    $  replicate depth ' '
    ++ (if collapsed n then "+" else " ")
    ++ label
    ++ replicate (width - depth - 1 - (length label)) ' '
  where
    label = labelFn $ value n


window start sz = take sz . drop start . concatMap (flatten . appendDepths 0 . filterVisible) . nodes

appendDepths d (Node n ns) = Node (d,n) $ map (appendDepths (d+1)) ns

filterVisible  (Node n ns) = Node n $ if collapsed n then [] else map filterVisible ns

